#include "ag.h"
#include "drivers/simple/mmc.h"
#include <rtl/print.h>
#include <rtl/list.h>
#include <rtl/checksum.h>
#include <rtl/heap.h>
#include <rtl/memory.h>
#include <rtl/string.h>

#define MODULE "PT"
#define PT_MAX_MMC 8
#define GPT_SIGNATURE 0x5452415020494645ULL

extern PVOID AgpGlobalHeap;
BOOLEAN AgpPartitionListInitialized = FALSE;
RTL_LIST AgpPartitionList[PT_MAX_MMC];

typedef struct _GPT_GUID
{
    UINT8 P0[4];
    UINT8 P1[2];
    UINT8 P2[2];
    UINT8 P3[2];
    UINT8 P4[6];
} GPT_GUID, *PGPT_GUID;

typedef struct _GPT_HEADER
{
    UINT64 Signature;
    UINT32 Revision;
    UINT32 HeaderSize;
    UINT32 HeaderCrc32;
    UINT32 Reserved;
    UINT64 CurrentLba;
    UINT64 BackupLba;
    UINT64 FirstUsableLba;
    UINT64 LastUsableLba;
    GPT_GUID DiskGuid;
    UINT64 PartitionTableStartLba;
    UINT32 PartitionEntriesCount;
    UINT32 PartitionEntrySize;
    UINT32 PartitionTableCrc32;
} PACKED GPT_HEADER, *PGPT_HEADER;

typedef struct _GPT_PTE
{
    GPT_GUID PartitionTypeGuid;
    GPT_GUID UniquePartitionGuid;
    UINT64 FirstLba;
    UINT64 LastLba;
    UINT64 Attributes;
    UINT16 Name[36];
} GPT_PTE, *PGPT_PTE;

VOID AgpPrintGptGuid(PGPT_GUID G, BOOLEAN NewLine)
{
    RtlPrintf("{%02X%02X%02X%02X-%02X%02X-%02X%02X-%02X%02X-%02X%02X%02X%02X%02X%02X}%s",
              G->P0[3],
              G->P0[2],
              G->P0[1],
              G->P0[0],
              G->P1[1],
              G->P1[0],
              G->P2[1],
              G->P2[0],
              G->P3[0],
              G->P3[1],
              G->P4[0],
              G->P4[1],
              G->P4[2],
              G->P4[3],
              G->P4[4],
              G->P4[5],
              NewLine ? "\n" : "");
}

BOOLEAN AgpVerifyGptHeader(PGPT_HEADER Header)
{
    UINT32 ComputedCrc;
    UINT32 T;

    if (Header->HeaderSize > MMC_BLOCK_SIZE)
    {
        RtlPrintf("[%s] Header size exceeds block size.\n", MODULE);
        return FALSE;
    }

    T = Header->HeaderCrc32;
    Header->HeaderCrc32 = 0;
    ComputedCrc = RtlComputeCrc32(0, Header, Header->HeaderSize);
    Header->HeaderCrc32 = T;

    if (ComputedCrc != Header->HeaderCrc32)
    {
        RtlPrintf("[%s] Header CRC32 mismatch, 0x%08x != 0x%08x.\n", MODULE, Header->HeaderCrc32, ComputedCrc);
        return FALSE;
    }

    return TRUE;
}

UINT64 AgpApplyPartitionPolicy(CONST CHAR *Name)
{
    if (RtlCompareString(Name, "boot", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "recovery", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "recoverbak", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "system", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "userdata", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "cache", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "cust", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "seccfg", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "logo", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else if (RtlCompareString(Name, "lk", ~0, TRUE) == 0)
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
    else
    {
        return AG_P_SECPOL_ALLOW_READ | AG_P_SECPOL_ALLOW_WRITE;
    }
}

VOID AgpProcessPte(INT Mmc, INT N, PGPT_PTE Pte)
{
    RtlPrintf("[%s] %5d 0x%016llx 0x%016llx 0x%016llx ",
              MODULE, N,
              Pte->FirstLba, Pte->LastLba,
              Pte->Attributes);
    AgpPrintGptGuid(&Pte->PartitionTypeGuid, FALSE);
    RtlPrintf(" ");
    AgpPrintGptGuid(&Pte->UniquePartitionGuid, FALSE);

    CHAR Name[36];

    // Simple UTF-16 to ASCII conversion, for now this should be enought
    INT I;
    for (I = 0; Pte->Name[I] != 0 && I < 36; I++)
    {
        if (Pte->Name[I] & 0xFF00)
            Name[I] = '?';
        else
            Name[I] = Pte->Name[I] & 0xFF;
    }
    Name[I] = '\0';
    RtlPrintf(" %s\n", Name);

    PAG_PTE AgPte = (PAG_PTE)RtlHeapAllocateMemory(AgpGlobalHeap, sizeof(AG_PTE));
    if (!AgPte)
    {
        RtlPrintf("[%s] Out of memory(?), dropping entry.\n");
        return;
    }
    RtlCopyMemory(AgPte->Name, Name, 36);
    AgPte->FirstLba = Pte->FirstLba;
    AgPte->LastLba = Pte->LastLba;
    AgPte->Policy = AgpApplyPartitionPolicy(AgPte->Name);
    RtlListAdd(&AgpPartitionList[Mmc], &AgPte->Node);
}

BOOLEAN AgpReadGptFrom(INT MmcControllerIndex, UINT32 Lba, BOOLEAN Primary)
{
    N1_STATUS Status;
    UINT8 Buffer[MMC_BLOCK_SIZE];
    Status = MtkMmcRead(MmcControllerIndex, Buffer, 1, Lba);
    if (!N1_SUCCESS(Status))
    {
        RtlPrintf("[%s] Failed to read partition table.\n", MODULE);
        return FALSE;
    }

    PGPT_HEADER GptHeader = (PGPT_HEADER)Buffer;
    if (GptHeader->Signature != GPT_SIGNATURE)
    {
        RtlPrintf("[%s] %s GPT not found.", MODULE, Primary ? "Primary" : "Secondary");
        return FALSE;
    }

    if (!AgpVerifyGptHeader(GptHeader))
    {
        RtlPrintf("[%s] %s GPT is corrupted.\n", MODULE, Primary ? "Primary" : "Secondary");
        return FALSE;
    }
    RtlPrintf("[%s] Disk GUID: ", MODULE);
    AgpPrintGptGuid(&GptHeader->DiskGuid, TRUE);
    // TODO: verify sanity of data stored in header eg.
    // check
    //   if partition table overlap with header or partitions
    //   if length of structures are ridiculously big or small (eg 0 bytes)
    //   if partition tables overlap with first usable sector
    //   etc.
    //   currently fedding agent with invalid data could result in completely undefined behaviour

    UINT32 PtFirstLba = GptHeader->PartitionTableStartLba;
    UINT32 PtEntrySize = GptHeader->PartitionEntrySize;
    UINT64 PtSize = GptHeader->PartitionEntrySize * GptHeader->PartitionEntriesCount;
    UINT64 PtNumBlocks = ALIGN_UP(PtSize, MMC_BLOCK_SIZE) / MMC_BLOCK_SIZE;
    UINT32 PtCrc32 = GptHeader->PartitionTableCrc32;

    RtlPrintf("[%s] %d partition entries of size %d bytes starting at LBA %d and occupying %lld blocks.\n",
              MODULE,
              GptHeader->PartitionEntriesCount,
              GptHeader->PartitionEntrySize,
              GptHeader->PartitionTableStartLba,
              PtNumBlocks);

    UINT32 PtComputedCrc32 = 0;
    UINT32 L = PtSize;

    for (UINT32 Lba = PtFirstLba; Lba - PtFirstLba < PtNumBlocks; Lba++)
    {
        if (!N1_SUCCESS(MtkMmcRead(MmcControllerIndex, Buffer, 1, Lba)))
        {
            RtlPrintf("[%s] Failed to read block %d\n", MODULE, Lba);
            return FALSE;
        }

        PtComputedCrc32 = RtlComputeCrc32(PtComputedCrc32, Buffer, MIN(L, MMC_BLOCK_SIZE));
        L -= MIN(L, MMC_BLOCK_SIZE);
    }

    if (PtComputedCrc32 != PtCrc32)
    {
        RtlPrintf("[%s] Partition table CRC32 mismatch, 0x%08x != 0x%08x.\n", MODULE, PtCrc32, PtComputedCrc32);
        return FALSE;
    }

    UINT PtEntriesPerBlock = MMC_BLOCK_SIZE / PtEntrySize;
    if (MMC_BLOCK_SIZE % PtEntrySize != 0)
    {
        RtlPrintf("[%s] Partition table entries cross block boundary.\n", MODULE);
        return FALSE;
    }

    INT J = 0;
    for (UINT32 Lba = PtFirstLba; Lba - PtFirstLba < PtNumBlocks; Lba++)
    {
        if (!N1_SUCCESS(MtkMmcRead(MmcControllerIndex, Buffer, 1, Lba)))
        {
            RtlPrintf("[%s] Failed to read block %d\n", MODULE, Lba);
            return FALSE;
        }

        for (UINT I = 0; I < PtEntriesPerBlock; I++)
        {
            PGPT_PTE Pte = (PGPT_PTE)((UINTPTR)Buffer + PtEntrySize * I);
            AgpProcessPte(MmcControllerIndex, J, Pte);
            J++;
        }
    }

    return TRUE;
}

BOOLEAN AgpReadPartitionTable(INT MmcControllerIndex)
{
    RtlPrintf("[%s] Reading partition table.\n", MODULE);

    if (!AgpPartitionListInitialized)
    {
        for (INT I = 0; I < PT_MAX_MMC; I++)
            RtlListInitialize(&AgpPartitionList[I]);
    }

    if (AgpReadGptFrom(MmcControllerIndex, 1, TRUE))
    {
        AgpPartitionListInitialized = TRUE;
        return TRUE;
    }

    return FALSE;
}

PAG_PTE AgpLookupPartition(INT Mmc, CONST CHAR *Name)
{
    if (!AgpPartitionListInitialized || Mmc >= PT_MAX_MMC)
        return NULL;

    PAG_PTE Pte;
    RtlListForEach(Pte, &AgpPartitionList[Mmc], Node)
    {
        if (RtlCompareString(Pte->Name, Name, ~0, TRUE) == 0)
            return Pte;
    }

    return NULL;
}